# Copyright (c) 2024-present AI-Labs

import streamlit as st
import os
import base64
import requests, json
from io import BytesIO
from PIL import Image
from datetime import datetime
from database.database import engine
from sqlalchemy import text

from common.draw import init_draw_config_form, save_text2image, save_image2image, save_inpainting
from utils import init_page_header, init_session_state, get_avatar
from utils import select_aigc_left_freq, update_aigc_perm_freq, use_limited, check_use_limit
from utils import main_container_height
from streamlit_drawable_canvas import st_canvas
 
title = "图片生成"
icon = "🎨"
# 初始化页面布局
init_page_header(title, icon)
# 初始化会话Session
init_session_state()


# 查询当天的AI使用次数，用于限制用户使用次数
def select_images_freq():
    count = 0
    try:
        with engine.connect() as conn:
            sql = text(f'''
            select count(*) from ai_labs_images where username = :username and date_time >= current_date;
            ''')
            count = conn.execute(sql, [{'username': st.session_state.username}]).fetchone()[0]
    except Exception as e:
        st.exception(e)
    return count


# 查询当前的可使用AI的永久次数、临时次数、已使用次数，用于限制用户使用次数
def select_aigc_freq():
    st.session_state.aigc_temp_freq, st.session_state.aigc_perm_freq = 100, 100 # select_aigc_left_freq()
    st.session_state.aigc_temp_images = 0 # select_images_freq()


# 插入一次绘图记录到数据库
def insert_images(user, assistant):
    try:
        with engine.connect() as conn:
            date_time = datetime.now()
            sql = text(f'''
            insert into ai_labs_images(username, user, assistant, date_time) values(:username, :user, :assistant, :date_time)
            ''')
            conn.execute(sql, [{
                'username': st.session_state.username,
                'user': user,
                'assistant': assistant,
                'date_time': date_time
            }])
            conn.commit()
    except Exception as e:
        st.error(e)


# 当发生图片优化、局部重绘后，用于更新绘图记录
def update_images(id, assistant):
    try:
        with engine.connect() as conn:
            date_time = datetime.now()
            sql = text(f'''
            update ai_labs_images set assistant = :assistant where id = :id
            ''')
            conn.execute(sql, [{
                'assistant': assistant,
                'id': id
            }])
            conn.commit()
    except Exception as e:
        st.error(e)


# 查询最后10条记录，用于页面展示
def select_images():
    with engine.connect() as conn:
        sql = text("""
            select id,user,assistant from (select * from ai_labs_images where username = :username order by id desc limit :showlimit) temp order by id
        """)
        return conn.execute(sql, [{
            'username': st.session_state.username,
            'showlimit': st.session_state.showlimit
        }]).fetchall()

# 图片生成的主要逻辑
def cache_draw(user_input_text):
    try:
        localfile = save_text2image(user_input_text)
        st.image(localfile)
        cols = st.columns([0.2,0.2,0.6])
        insert_images(user_input_text, localfile)
        st.rerun()
    except Exception as ex:
        st.exception(ex)


# 将图片转换成Base64编码的字符串
def image_to_base64(pil_image: Image)->str:
    buffered = BytesIO()
    pil_image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode('utf-8')


# 图片优化（图生图）功能的界面、逻辑
@st.experimental_dialog("📝图片优化")
def image2image(id, assistant):
    init_image = st.image(assistant)
    prompt = st.chat_input("请告诉我如何优化")
    if prompt:
        init_image = Image.open(assistant)
        try:
            localfile = save_image2image(prompt, f"data:image;base64,{image_to_base64(init_image)}")
            st.session_state.image2image = localfile
            st.image(localfile)
        except Exception as ex:
            st.exception(ex)
    if "image2image" in st.session_state.keys():
        cols = st.columns(2)
        with cols[0]:
            if st.button("放弃", use_container_width=True):
                del st.session_state.image2image
                st.rerun()
        with cols[1]:
            if st.button("接受", use_container_width=True, type="primary"):
                update_images(id, st.session_state.image2image)
                del st.session_state.image2image
                st.rerun()
    else:
        if st.button("关闭", use_container_width=True):
            st.rerun()


# 局部重绘功能的界面、逻辑
@st.experimental_dialog("📝局部重绘")
def inpainting(id, assistant):
    background_image = Image.open(assistant)
    canvas = st_canvas(key=f"canvas_{id}",
            background_image=background_image,
            width=450,
            height=450,
            stroke_width=20,
            stroke_color="#CCCCCC",
            )
    prompt = st.chat_input("请告诉我如何重绘")
    if prompt:
        init_image = Image.open(assistant)
        mask_image = Image.fromarray(canvas.image_data).resize(background_image.size).convert("L")
        try:
            localfile = save_inpainting(prompt, f"data:image;base64,{image_to_base64(init_image)}", f"data:image;base64,{image_to_base64(mask_image)}")
            st.session_state.inpainting = localfile
            st.image(localfile)
        except Exception as ex:
            st.exception(ex)
    if "inpainting" in st.session_state.keys():
        cols = st.columns(2)
        with cols[0]:
            if st.button("放弃", use_container_width=True):
                del st.session_state.inpainting
                st.rerun()
        with cols[1]:
            if st.button("接受", use_container_width=True, type="primary"):
                update_images(id, st.session_state.inpainting)
                del st.session_state.inpainting
                st.rerun()
    else:
        if st.button("关闭", use_container_width=True):
            st.rerun()

if __name__ == '__main__':
    
    select_aigc_freq()

    with st.sidebar:
        tabs = st.tabs(["图片设置"])
        with tabs[0]:
            init_draw_config_form()
        if check_use_limit:
            st.info(f"免费次数已用：{min(st.session_state.aigc_temp_freq, st.session_state.aigc_temp_images)}/{st.session_state.aigc_temp_freq} 次。\n\r永久次数剩余：{st.session_state.aigc_perm_freq} 次。", icon="🙂")

    main_container = st.container(height=main_container_height)

    for id, user, assistant in select_images():
        with main_container.chat_message("user"):
            st.write(user)
        with main_container.chat_message("assistant", avatar=get_avatar("stable-diffusion")):
            st.image(assistant)
            cols = st.columns([0.2,0.2,0.6])
            with cols[0]:
                if st.button("📝图片优化", key=f"image2image_{id}"):
                    image2image(id, assistant)
            with cols[1]:
                if st.button("📝局部重绘", key=f"inpainting_{id}"):
                    inpainting(id, assistant)
                # with st.popover("📝图片优化"):
                #     with st.form(f"edit_image_{id}_form_0"):
                #         init_image = st.image(assistant)
                #         prompt = st.text_input(label="", placeholder="请输入优化方案")
                #         if st.form_submit_button("提交", type="primary", use_container_width=True):
                #             pass
            

    user_input_text = st.chat_input("您的输入...")

    if user_input_text:
        with main_container.chat_message("user"):
            st.write(user_input_text)
        
        with main_container.chat_message("assistant", avatar=get_avatar("stable-diffusion")):
            if check_use_limit and st.session_state.aigc_temp_images >= st.session_state.aigc_temp_freq and st.session_state.aigc_perm_freq < 1:
                use_limited()
            else:
                with st.spinner("处理中，请稍等..."):
                    # st.image("statics/images/loading.gif")
                    cache_draw(user_input_text)
                    if check_use_limit and st.session_state.aigc_temp_images >= st.session_state.aigc_temp_freq:
                        update_aigc_perm_freq(-1)
                    select_aigc_freq()
